import time
from kazoo.client import KazooClient
from logger_config import logger_factory

logger = logger_factory()


class ZKController:
    def __init__(self, zk_hosts, election_path, session_timeout, identifier):
        self.zk = KazooClient(hosts=zk_hosts, timeout=session_timeout)
        self.election_path = election_path
        self.identifier = identifier
        self.leader = None
        self.is_leader = False

    def start(self):
        self.zk.start()
        self.zk.ensure_path(self.election_path)
        print(self.election_path)

    def run_func(self, func):
        self._run_election(func)

    def _run_election(self, func):
        try:
            # 创建顺序节点
            path = self.zk.create(self.election_path + '/' + self.identifier[:-1], ephemeral=True, sequence=True)
            node_name = path.split('/')[-1]
            # 等待节点变化
            while True:
                children_list = self.zk.get_children(self.election_path)
                children_list.sort()
                current_index = children_list.index(node_name)
                if current_index == 0:
                    if not self.is_leader:
                        logger.info(f'{self.identifier} become queen leader')
                        self.is_leader = True
                        self._become_leader(func)
                else:
                    if self.is_leader:
                        logger.info(f'{self.identifier} not as leader')
                        self.is_leader = False
                        self._close_leader()
        except Exception as e:
            logger.error(f'error occur iin queen leader election {e}')
        finally:
            self.zk.stop()
            self.zk.close()

    def _become_leader(self, func):
        # 实现领导者逻辑
        while self.is_leader:
            logger.info(f'{self.identifier} 正在服务……')
            func()

    def _close_leader(self):
        # 实现非领导者逻辑
        logger.info(f'{self.identifier} 正在等待成为领导者')
        time.sleep(60)